package com.forever.controller;

import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;

import java.time.Duration;
import java.util.List;

import static dev.langchain4j.agent.tool.JsonSchemaProperty.INTEGER;
import static dev.langchain4j.data.message.UserMessage.userMessage;
import static java.util.Collections.singletonList;

public class OpenAiOllamaTool {
    public static void main(String[] args) {
        ToolSpecification calculator = ToolSpecification.builder()
                .name("add")
                .description("计算两个数的和")
                .addParameter("第一个数的值", INTEGER)
                .addParameter("第二个数的值", INTEGER)
                .build();
        List<ToolSpecification> toolSpecifications = singletonList(calculator);

        Duration duration = Duration.ofSeconds(200);
        ChatLanguageModel model =
                new OpenAiChatModel
                        .OpenAiChatModelBuilder()
                        .baseUrl("http://localhost:11434/v1")
                        .apiKey("ollama")
                        .modelName("qwen2:7b")
                        .logRequests(true)
                        .logResponses(true)
                        .timeout(duration)
                        .build();
        UserMessage userMessage = userMessage("计算3加5的值");
        Response<AiMessage> response = model.generate(singletonList(userMessage),toolSpecifications);
        System.out.println(response);
    }
}
